import datetime
import gradio as gr
import pymongo as pm
from util_mongo import enqueue, get_sorted_by_key
from common import MONGODB_NAME, VALUE, KEY, IO_PREFIX
from util import merge_dialog_in_and_out, translate_ns
import time
from transformers import AutoTokenizer, AutoModel

# 连接Mongodb
mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
mdb = mongo[MONGODB_NAME]

# 目前暂时不区分用户，统一用u0001
username = 'u0001'

print('-------------------------------------------------------')
print('正在加载模型……')
# model_name = "THUDM/chatglm2-6b-int4"
model_name = "/root/.cache/huggingface/hub/models--THUDM--chatglm2-6b-int4/snapshots/66ecaf1db3a5085714e133357ea4824b69698743"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name,trust_remote_code=True).cuda()
model = model.eval()
print('模型已经加载完毕。')


def chat_once(xinput):
    """输入文本，用假模型返回输出。"""
    # 获取聊天历史
    xlog = get_history()
    
    # 得到输出
    xoutput, history = model.chat(tokenizer, xinput, history=xlog)
    
    # 输入存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_in', username, ts, {
        VALUE: xinput,
        'io': 'i',
    })
    
    # 输出存入mongodb
    ts = time.time_ns()
    enqueue(mdb, 'dialog_out', username, ts, {
        VALUE: xoutput,
        'io': 'o',
    })

    # 获取聊天历史
    xlog = get_history()

    # 返回空字符串清空输入文本框，返回聊天历史
    return '', xlog


def get_history():
    """
    获取聊天历史

    把历史组合成对子的列表，例如：
    in: xxx
    out: yyy
    in: zzzz
    in: aaaa
    out: bbbb
    out: kkkk
    =>
    [
        [xxx, yyy],
        [zzzz, None],
        [aaaa, bbbb],
        [None, kkkk],
    ]

    """
    # 获取输入历史
    rows_in = get_sorted_by_key(mdb, 'dialog_in', username, limit=6, is_keep_others=False)
    # 获取输出历史
    rows_out = get_sorted_by_key(mdb, 'dialog_out', username, limit=6, is_keep_others=False)
    # 合并输入历史和输出历史
    rows = merge_dialog_in_and_out(rows_in, rows_out)  # DESC
    # 按时间正序排列
    rows = rows[::-1]  # ASC

    if not len(rows):
        # 没有历史则返回空列表
        log = []
    else:
        # 把聊天历史组合成对子的列表
        log = []
        pair = [None, None]
        for row in rows:
            text = row[VALUE]
            timestamp = row[KEY]
            dt_str = translate_ns(timestamp)
            xstr = IO_PREFIX.get(row["io"], '?') + ': '
            xstr += f'[ {dt_str} ] '

            xstr += text

            if 'i' == row['io']:
                # 如果是输入

                # 如果对子里面已经有输入或输出了，则新建对子
                if pair[0] is not None or pair[1] is not None:
                    log.append(pair)
                    pair = [xstr, None]
                    continue

                # 输入放入对子的第一个元素
                pair[0] = xstr

            elif 'o' == row['io']:
                # 如果是输出

                # 如果对子里面已经有输出了，则新建对子
                if pair[1] is not None:
                    log.append(pair)
                    pair = [None, xstr]
                    continue

                # 输出放入对子的第2个元素
                pair[1] = xstr
        log.append(pair)

    return log


# 构建Gradio界面
with gr.Blocks(analytics_enabled=False) as demo:

    # 输入文本框
    cmp_input = gr.Textbox(interactive=True, label='输入')

    # 输入按钮
    cmp_input_button = gr.Button('发送', )

    # 聊天历史
    cmp_chatbot = gr.Chatbot().style(height=350)

    # 输入按钮点击事件
    cmp_input_button.click(chat_once, [cmp_input, ], [
        cmp_input,
        cmp_chatbot,
    ], queue=False)

    # 加载时展示聊天历史
    demo.load(get_history, None, [
        cmp_chatbot,
    ], queue=False)

# 使用queue（目前这个例子太简单，没有用到queue，后面会用到。）
demo.queue(
    concurrency_count=10,
    status_update_rate='auto',
    # status_update_rate=0.02,
)

# 启动Gradio
# demo.launch(server_name='0.0.0.0', server_port=6006, share=True, debug=True)  # 带内网穿透
demo.launch(server_name='0.0.0.0', server_port=6006, debug=True)
